# storage/buffer.py

"""
实现了缓冲池管理器和页面替换策略。
"""
from typing import Dict, List, Deque
from collections import OrderedDict
from collections import deque
# import logging
from readerwriterlock import rwlock
import logging
import time

# 导入 Page 类，它定义了页面的内存结构
from .page import Page

# 设置日志，方便调试和展示
# logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s: %(message)s')

class ReplacementPolicy:
    
    """
    抽象的页面替换策略基类。
    定义了所有替换策略必须实现的接口。
    """
    def __init__(self, pool_size: int, buffer_manager):
        self.pool_size = pool_size
        self.buffer_manager = buffer_manager  # 存储 BufferManager 的引用，用于访问缓冲池

    def page_accessed(self, page_id: int):
        """当一个页面被访问时调用。"""
        raise NotImplementedError

    def get_victim_page(self) -> int:
        """选择并返回一个将被淘汰的页面ID。"""
        raise NotImplementedError

    def add_page(self, page_id: int):
        """当新页面被加入缓冲池时调用。"""
        raise NotImplementedError

    def remove_page(self, page_id: int):
        """当页面从缓冲池中移除时调用。"""
        raise NotImplementedError

class LRUReplacementPolicy(ReplacementPolicy):
    
    """
    最近最少使用 (LRU) 替换策略，使用 OrderedDict 实现。
    """
    def __init__(self, pool_size: int, buffer_manager):
        super().__init__(pool_size, buffer_manager)
        self.lru_cache: OrderedDict[int, None] = OrderedDict()

    def page_accessed(self, page_id: int):
        """
        当页面被访问时，将其移动到 OrderedDict 的末尾。
        """
        if page_id in self.lru_cache:
            # 移除页面
            self.lru_cache.pop(page_id)
        # 将其重新添加到末尾，使其成为最新使用的
        self.lru_cache[page_id] = None

    def get_victim_page(self) -> int:
        """
        返回LRU队列中可被淘汰的（pin_count=0）最久未使用的页面ID。
        """
        for page_id in self.lru_cache:
            page = self.buffer_manager.buffer_pool.get(page_id)
            if page and page.pin_count == 0:
                # 找到了可淘汰的页面，从缓存中移除
                self.lru_cache.pop(page_id)
                return page_id
        
        raise RuntimeError("Buffer pool is full and all pages are pinned.")

    def add_page(self, page_id: int):
        """新页面加入时，直接添加到 OrderedDict 末尾。"""
        self.lru_cache[page_id] = None

    def remove_page(self, page_id: int):
        """从缓存中移除页面。"""
        if page_id in self.lru_cache:
            self.lru_cache.pop(page_id)

class FIFOReplacementPolicy(ReplacementPolicy):
    
    """
    先进先出 (FIFO) 替换策略。
    """
    def __init__(self, pool_size: int, buffer_manager):
        super().__init__(pool_size, buffer_manager)
        self.queue: Deque[int] = deque()

    def page_accessed(self, page_id: int):
        # FIFO策略不关心访问顺序，所以该方法不做任何事
        pass

    def get_victim_page(self) -> int:
        """
        返回 FIFO 队列中可被淘汰的（pin_count=0）最旧的页面。
        """
        for _ in range(len(self.queue)):
            page_id = self.queue.popleft()
            page = self.buffer_manager.buffer_pool.get(page_id)
            if page and page.pin_count == 0:
                return page_id
            else:
                self.queue.append(page_id)

        raise RuntimeError("Buffer pool is full and all pages are pinned.")

    def add_page(self, page_id: int):
        self.queue.append(page_id)

    def remove_page(self, page_id: int):
        if page_id in self.queue:
            self.queue.remove(page_id)

class BufferManager:
    """
    管理缓冲池中的数据页面。
    它负责处理页面缓存、替换和I/O操作。
    """
    def __init__(self, pool_size: int, file_manager, replacement_policy_class=LRUReplacementPolicy):
        self.pool_size = pool_size
        self.file_manager = file_manager
        self.buffer_pool: Dict[int, Page] = {}
        self.hit_count = 0
        self.miss_count = 0
        self.replacement_events: List[dict] = []  # 记录淘汰日志
        self.policy = replacement_policy_class(pool_size, self)
        # 这个锁只用于保护 buffer_pool 字典本身，而不是页面内容或元数据
        self.rw_lock = rwlock.RWLockFair()

    # _pin_page 和 _unpin_page 方法不再需要，因为我们现在使用细粒度锁
    
    def new_page(self) -> Page:
        """
        创建并返回一个新的页面。
        这个操作会改变缓冲池的结构，所以需要写锁。
        """
        with self.rw_lock.gen_wlock():
            page_id = self.file_manager.allocate_page()
           # logging.info(f"Allocating new page with id: {page_id}")
            
            if len(self.buffer_pool) >= self.pool_size:
                self._evict_page()

            new_page = Page(page_id)
            with new_page.lock:
                new_page.pin_count += 1
            self.buffer_pool[page_id] = new_page
            self.policy.add_page(page_id)
            
            return new_page

    def get_page(self, page_id: int) -> Page:
        """
        获取一个页面，使用细粒度锁优化并发读性能。
        """
        # 阶段 1：使用读锁快速检查缓存
        with self.rw_lock.gen_rlock():
            if page_id in self.buffer_pool:
                self.hit_count += 1
                page = self.buffer_pool[page_id]
                
                # 在获取页面的细粒度锁后，再更新pin_count和访问策略
                with page.lock:
                    page.pin_count += 1
                    self.policy.page_accessed(page_id)
                   # logging.info(f"Pinned page {page_id} with fine-grained lock. New pin count: {page.pin_count}")
                
                return page

        # 阶段 2：如果缓存未命中，获取写锁来加载页面
        with self.rw_lock.gen_wlock():
            # 再次检查，防止在释放读锁到获取写锁期间，其他线程已经加载了页面
            if page_id in self.buffer_pool:
                self.hit_count += 1
                page = self.buffer_pool[page_id]
                with page.lock:
                    page.pin_count += 1
                    self.policy.page_accessed(page_id)
                return page
            
            # 确定页面确实不在缓存中，执行加载和淘汰
            self.miss_count += 1
        #    logging.warning(f"Cache Miss for page {page_id}. Fetching from disk...")

            if len(self.buffer_pool) >= self.pool_size:
                self._evict_page()

            page_data = self.file_manager.read_page(page_id)
            if not page_data: 
                raise ValueError(f"Page {page_id} not found on disk.")
            
            page = Page(page_id, page_data)
            with page.lock:
                page.pin_count += 1
            self.buffer_pool[page_id] = page
            self.policy.add_page(page_id)
            
            return page

    def unpin_page(self, page_id: int, is_dirty: bool = False):
        """
        释放页面。使用页面自己的锁来修改状态。
        """
        with self.rw_lock.gen_rlock():  # 只需要读锁来查找页面
            if page_id not in self.buffer_pool:
                return
            page = self.buffer_pool[page_id]
        
        with page.lock:  # 使用页面的细粒度锁来修改状态
            if page.pin_count > 0:
                page.pin_count -= 1
            if is_dirty:
                page.is_dirty = True
            #  logging.info(f"Unpinned page {page_id}. Pin count: {page.pin_count}")

    def flush_page(self, page_id: int):
        """
        强制将一个页面写回磁盘。
        使用页面的细粒度锁来保护操作。
        """
        with self.rw_lock.gen_rlock():
            page = self.buffer_pool.get(page_id)

        if page:
            with page.lock:
                if page.is_dirty:
                    self.file_manager.write_page(page.page_id, page.data)
                    page.is_dirty = False
                    # logging.info(f"Forced flush of dirty page {page_id}.")
                # else:
                    # logging.info(f"Page {page_id} is not dirty, no need to flush.")

    def flush_all_pages(self):
        """
        将缓冲池中所有的脏页都写回磁盘。
        需要写锁，因为会迭代并可能从字典中删除页面。
        """
        with self.rw_lock.gen_wlock():
            # print("Flushing all dirty pages from buffer pool...")
            for page_id in list(self.buffer_pool.keys()):
                page = self.buffer_pool[page_id]
                with page.lock:
                    if page.is_dirty:
                        self.file_manager.write_page(page.page_id, page.data)
                        page.is_dirty = False
                        # logging.info(f"Forced flush of dirty page {page_id}.")
    
    def _evict_page(self):
        """
        根据替换策略选择一个页面进行淘汰。
        这个方法在 new_page 和 get_page 中调用，它们已经持有写锁。
        """
        # 策略查找受限于 BufferManager 的 rw_lock
        victim_id = self.policy.get_victim_page()
        victim_page = self.buffer_pool.get(victim_id)

        # 在执行淘汰前，需要锁定该页面，以确保其状态不被其他线程修改
        if victim_page:
            with victim_page.lock:
                if victim_page.pin_count > 0:
                    # 再次检查，以防在获取锁前状态发生变化
                    # logging.error(f"Attempting to evict pinned page {victim_id}. This is a bug in replacement policy.")
                    raise RuntimeError("Replacement policy attempted to evict a pinned page.")

                # 执行淘汰操作
                del self.buffer_pool[victim_id]
                self.policy.remove_page(victim_id)

                if victim_page.is_dirty:
                    self.file_manager.write_page(victim_page.page_id, victim_page.data)
                    logging.info(f"Evicted and flushed dirty page {victim_id}.")
                else:
                    logging.info(f"Evicted clean page {victim_id}.")
                # 记录淘汰事件
                self.replacement_events.append({
                    'victim': victim_id,
                    'dirty': victim_page.is_dirty,
                    'timestamp': time.time()
                })
                
                return victim_id
        else:
            raise RuntimeError(f"Replacement policy returned non-existent page {victim_id}")

    def mark_dirty(self, page_id: int):
        """Utility: mark a page as dirty without changing pin count."""
        with self.rw_lock.gen_rlock():
            page = self.buffer_pool.get(page_id)
        if page:
            with page.lock:
                page.is_dirty = True

    def get_stats(self) -> Dict[str, float]:
        """
        返回缓存命中率统计。
        只读操作，使用读锁。
        """
        with self.rw_lock.gen_rlock():
            total_accesses = self.hit_count + self.miss_count
            hit_ratio = (self.hit_count / total_accesses) * 100 if total_accesses > 0 else 0
            return {
                "total_accesses": total_accesses,
                "hit_count": self.hit_count,
                "miss_count": self.miss_count,
                "hit_ratio_percent": round(hit_ratio, 2)
            }